/*
 * Compute corrcoef matrix for 2 inputs A and B. 
 * Assuming that A and B must have the same size.
 */
#include <stdio.h>
#include <stdlib.h>

#include "util.h"
#include "common.h"
#include "corrcoefmatrix_kernel.h"
#include "sum_sum2_kernel.h"

void usage(int argc, char **argv) {
   fprintf(stderr,"Usage: %s <filename> <filename>\n",argv[0]);
}

int main(int argc, char **argv) {
  char *ctrl_filename; 
  char *case_filename;
  float *inA, *inB, *out;
  int rowA,colA;
  int rowB,colB;
  int padded_row;
  int grid_size_x,grid_size_y;
  int sizeIn, sizeOut;

  cudaError_t err;
  float *inAG, *inBG, *outG;

  double st,et;
  
  if (argc < 3) {
     usage(argc,argv);
     exit(1);
  }

  ctrl_filename = argv[1];
  inA = read_matrix(ctrl_filename,TILESIZE,&rowA,&colA,&padded_row); 
  sizeIn = padded_row*colA*sizeof(float);
  sizeOut = padded_row*padded_row*sizeof(float);
  out = (float *) malloc(sizeOut);
  
  case_filename = argv[2];
  inB = read_matrix(case_filename,TILESIZE,&rowB,&colB,&padded_row); 
  if (rowA != rowB) {
     fprintf(stderr,"# of rows in A and B are not equal\n");
     exit(1);
  }
  if (colA != colB) {
     fprintf(stderr,"# of columns in A and B are not equal\n");
     exit(1);
  }

  st = getTimeStamp();

  fprintf(stderr,"Running in CUDA mode...\n");
  dim3 blockDim(TILESIZE,TILESIZE);
  grid_size_x = padded_row/blockDim.x;
  grid_size_y = padded_row/blockDim.y;
  dim3 gridDim(grid_size_x,grid_size_y);
  
  if ( (grid_size_x > 65535) || (grid_size_y > 65535) ) {
     fprintf(stderr,"Grid size too large %dx%d\n",grid_size_x,grid_size_y);
     exit(1);
  } 
  fprintf(stderr,"Grid size %dx%d\n",grid_size_x,grid_size_y);

  if ((err=cudaMalloc((void **) &inAG,sizeIn)) != cudaSuccess) {
     fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
     exit(1);
  }
  if ((err=cudaMalloc((void **) &inBG,sizeIn)) != cudaSuccess) {
     fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
     exit(1);
  }
  if ((err=cudaMalloc((void **) &outG,sizeOut)) != cudaSuccess) {
     fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
     exit(1);
  }
  if ((err=cudaMemcpy(inAG,inA,sizeIn,cudaMemcpyHostToDevice)) != cudaSuccess) {
     fprintf(stderr,"cudaMemcpy from host to device fail: %s\n",cudaGetErrorString(err));
     exit(1);
  }
  if ((err=cudaMemcpy(inBG,inB,sizeIn,cudaMemcpyHostToDevice)) != cudaSuccess) {
     fprintf(stderr,"cudaMemcpy from host to device fail: %s\n",cudaGetErrorString(err));
     exit(1);
  }

  fprintf(stderr,"Precompute sum and sum^2 in CUDA...\n");
  float *sumAG, *sum2AG;
  float *sumBG, *sum2BG;

  if ((err=cudaMalloc((void **) &sumAG,padded_row*sizeof(float))) != cudaSuccess) {
     fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
     exit(1);
  }
  if ((err=cudaMalloc((void **) &sum2AG,padded_row*sizeof(float))) != cudaSuccess) {
     fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
     exit(1);
  }
  if ((err=cudaMalloc((void **) &sumBG,padded_row*sizeof(float))) != cudaSuccess) {
     fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
     exit(1);
  }
  if ((err=cudaMalloc((void **) &sum2BG,padded_row*sizeof(float))) != cudaSuccess) {
     fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
     exit(1);
  }
  dim3 blockDim1(TILESIZE,TILESIZE); 
  dim3 gridDim1(1,padded_row/TILESIZE);
  sum_sum2_kernel<<<gridDim1,blockDim1>>>(sumAG,sum2AG,inAG,padded_row,colA); 
  cudaThreadSynchronize();
  sum_sum2_kernel<<<gridDim1,blockDim1>>>(sumBG,sum2BG,inBG,padded_row,colB);
  cudaThreadSynchronize();
  corrcoefmatrixAB_kernel<<<gridDim,blockDim>>>(outG,inAG,inBG,sumAG,sum2AG,sumBG,sum2BG,padded_row,colA);  
  cudaFree(sumAG);
  cudaFree(sum2AG);
  cudaFree(sumBG);
  cudaFree(sum2BG);

  cudaThreadSynchronize();

  // check for error
  check_cuda_errors(__FILE__,__LINE__);

  cudaMemcpy(out,outG,sizeOut,cudaMemcpyDeviceToHost);
  cudaFree(inAG);
  cudaFree(inBG);
  cudaFree(outG);

  et = getTimeStamp();
  write_submatrix(out,padded_row,padded_row,rowA,rowA);

  fprintf(stderr,"Calculation time = %d ms.\n",(int) (et-st));
}


